import SimpleITK as sitk 
import numpy as np 
import pydicom
import pydicom_seg
import os 

def save_segmentation_to_dicom(t1_dir, segmentation_data, study_id):
    template = pydicom_seg.template.from_dcmqi_metainfo('metainfo.json')

    writer = pydicom_seg.MultiClassWriter(
        template=template,
        inplane_cropping=False,  # Crop image slices to the minimum bounding box on
                                # x and y axes
        skip_empty_slices=False,  # Don't encode slices with only zeros
        skip_missing_segment=False,  # If a segment definition is missing in the
                                    # template, then raise an error instead of
                                    # skipping it.
    )

    reader = sitk.ImageSeriesReader()
    series_IDs = reader.GetGDCMSeriesIDs(t1_dir)
    series = series_IDs[0]
    dcm_files = reader.GetGDCMSeriesFileNames(t1_dir, series)
    reader.SetFileNames(dcm_files)
    image = reader.Execute()
    image_data = sitk.GetArrayFromImage(image)

    print("image_data.shape",image_data.shape)  # 320

    # segmentation_data = np.zeros_like(image_data).astype(np.uint8)

    # segmentation_data[10:200] = 1
    segmentation_data = segmentation_data.astype(np.uint8)
    segmentation = sitk.GetImageFromArray(segmentation_data)
    segmentation.CopyInformation(image)

    source_images = [
        pydicom.dcmread(x, stop_before_pixels=False)
        for x in dcm_files
    ]
    dcm = writer.write(segmentation, source_images)

    os.makedirs("../segmentation_saved", exist_ok=True)

    dcm.save_as(f'../segmentation_saved/{study_id}_segmentation.dcm')
